import json
import pandas as pd

class MemoryHandler:
    def __init__(self, memory_file="memory.json"):
        self.memory_file = memory_file
        self.memory = self.load_memory()

    def load_memory(self):
        """从文件加载记忆数据。"""
        try:
            with open(self.memory_file, 'r') as file:
                return json.load(file)
        except FileNotFoundError:
            return {}

    def save_memory(self):
        """将记忆数据保存到文件。"""
        with open(self.memory_file, 'w') as file:
            json.dump(self.memory, file)

    def update_memory(self, key, value):
        """更新记忆中的数据。如果是 DataFrame，将其转换为字典格式。"""
        if isinstance(value, pd.DataFrame):
            # 将 DataFrame 转换为字典格式
            self.memory[key] = value.to_dict(orient='list')  # 使用 'list' 格式，列名作为键，列值作为列表
        elif isinstance(value, pd.Series):
            # 如果是 Series 也转换为字典
            self.memory[key] = value.to_dict()
        else:
            self.memory[key] = value
        self.save_memory()

    def get_memory(self, key):
        """检索记忆中的数据。"""
        return self.memory.get(key, None)


